# -*- coding:utf-8 -*-
# Author:      Boll.lai
import ldap3
from ldap3.core.exceptions import LDAPBindError
from ldap3.extend.standard.modifyPassword import ModifyPassword
from ldap3.utils.ciDict import CaseInsensitiveDict
from passlib.hash import ldap_salted_sha1

from common.utils import to_str
from setting import config

__all__ = ["AuthLdapHandler"]


class AuthLdapHandler(object):
    # def __init__(self, config):
    def __init__(self, LDAP_HOST=config.LDAP_HOST, LDAP_PORT=config.LDAP_PORT, LDAP_PROTOCOL=config.LDAP_PROTOCOL,
                 LDAP_USER=config.LDAP_USER, LDAP_AUTH=config.LDAP_AUTH, LDAP_BASE_DN=config.LDAP_BASE_DN, LDAP_USER_DN=config.LDAP_USER_DN,
                 LDAP_GROUP_DN=config.LDAP_GROUP_DN):
        # print(LDAP_HOST)
        # print(LDAP_PORT)
        # print(LDAP_USER)
        # print(LDAP_AUTH)
        # print(LDAP_PROTOCOL)
        # print(LDAP_BASE_DN)
        # print(LDAP_USER_DN)
        # print(LDAP_GROUP_DN)
        # self.host = config.get("LDAP_HOST")
        # self.port = config.get("LDAP_PORT")
        # self.user = config.get("LDAP_USER")
        # self.auth = config.get("LDAP_AUTH")
        # self.protocol = config.get("LDAP_PROTOCOL")
        # self.base_dn = config.get("LDAP_BASE_DN")
        # self.user_dn = config.get("LDAP_USER_DN")
        # self.group_dn = config.get("LDAP_GROUP_DN")
        self.host = LDAP_HOST
        self.port = LDAP_PORT
        self.user = LDAP_USER
        self.auth = LDAP_AUTH
        self.protocol = LDAP_PROTOCOL
        self.base_dn = LDAP_BASE_DN
        self.user_dn = LDAP_USER_DN
        self.group_dn = LDAP_GROUP_DN

    def bind(self):
        conn = ldap3.Connection(server=ldap3.Server(f'{self.protocol}://{self.host}:{self.port}'),
                                user=f'cn={self.user},{self.base_dn}', password=self.auth,
                                authentication=ldap3.SIMPLE)

        conn.bind()
        return conn if conn.bound else None

    def bind_user(self, uid, password, get_connection=False):
        user_dn = self.get_user_dn(uid)
        if not user_dn: return None

        conn = ldap3.Connection(
            server=ldap3.Server(f'{self.protocol}://{self.host}:{self.port}'),
            user=user_dn,
            password=str(password),
            authentication=ldap3.SIMPLE,
        )
        conn.bind()

        if get_connection:
            return conn

        bound = conn.bound
        conn.unbind()

        return bound

    def get_user_dn(self, uid):
        attributes = ["uid"]
        filter_syntax = "(&(uid=%s)(objectClass=posixAccount))" % uid

        conn = self.bind()
        conn.search(search_base=self.user_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)
        try:
            response = conn.response[0]
        except IndexError:
            return None

        conn.unbind()
        return response["dn"]

    def get_user_uidNumber(self, uid):
        attributes = ["uidNumber"]
        filter_syntax = "(&(uid=%s)(objectClass=posixAccount))" % uid

        conn = self.bind()
        conn.search(search_base=self.user_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)
        try:
            response = conn.response
        except IndexError:
            return None

        conn.unbind()
        return response[0]['attributes']['uidNumber']

    def get_group_dn(self, gid):
        attributes = ["cn"]
        filter_syntax = "(&(gidNumber=%s)(objectClass=posixGroup))" % gid

        conn = self.bind()
        conn.search(search_base=self.base_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)
        try:
            response = conn.response[0]
        except IndexError:
            return None

        conn.unbind()
        return response["dn"]

    def get_group_of_unique_dn(self, cn):
        attributes = ["cn"]
        filter_syntax = "(&(cn=%s)(objectClass=groupOfUniqueNames))" % cn

        conn = self.bind()
        conn.search(search_base=self.base_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)

        try:
            response = conn.response[0]
        except IndexError:
            return None

        conn.unbind()
        return response["dn"]

    def get_available_uid(self, gid):
        attributes = ['uidNumber']

        filter_syntax = "(&(gidNumber=%s)(objectClass=posixAccount))" % gid

        conn = self.bind()
        conn.search(search_base=self.user_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)

        uid_list = [
            entry["attributes"].get("uidNumber", -1)
            for entry in conn.response if isinstance(entry["attributes"], CaseInsensitiveDict)
        ]

        from random import randint
        return sorted(uid_list or [randint(10000, 50000)])[-1] + 1

    def get_ou_list(self, filter_syntax="(&(ou=*)(objectClass=organizationalUnit))"):
        attributes = ['ou']

        conn = self.bind()
        conn.search(search_base=self.base_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)

        ous = [
            {"name": entry['attributes']['ou'][0]}
            for entry in conn.response
            if 'attributes' in entry and all(attr in entry['attributes'] for attr in ['ou'])
        ]

        conn.unbind()
        return ous

    def get_groups(self, filter_syntax="(&(cn=*)(objectClass=posixGroup))"):
        attributes = ['cn', 'gidNumber', 'description']

        conn = self.bind()
        conn.search(search_base=self.group_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)

        groups = [
            {
                "id": entry['attributes']['gidNumber'],
                "name": entry['attributes']['cn'][0],
                "description": ''.join(entry['attributes']['description']),
            }
            for entry in conn.response
            if 'attributes' in entry and all(attr in entry['attributes'] for attr in attributes)
        ]

        conn.unbind()
        return groups

    def get_group_of_unique_list(self, filter_syntax="(&(cn=*)(objectClass=groupOfUniqueNames))"):
        attributes = ['cn', 'description']

        conn = self.bind()
        conn.search(search_base=self.group_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)

        for entry in conn.response:
            try:
                yield {
                    "description": entry['attributes']['description'][0],
                    "cn": entry['attributes']['cn'][0],
                }
            except IndexError:
                continue

        conn.unbind()

    def get_user(self, uid):
        attributes = [
            'uid',
            'userPassword',
            'description',
            'mail',
            'gidNumber',
            'telephoneNumber'
            # 'pwdAccountLockedTime',
        ]

        filter_syntax = "(&(uid=%s)(objectClass=posixAccount))" % uid

        conn = self.bind()
        conn.search(search_base=self.base_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)
        try:
            response = conn.response[0]
        except IndexError:
            return {}

        conn.unbind()
        # vpn_groups = [
        #     vpn_group["cn"]
        #     for vpn_group in self.get_group_of_unique_list(
        #         filter_syntax="(&(cn=*)(objectClass=groupOfUniqueNames)(uniqueMember=%s))" % self.get_user_dn(uid)
        #     )
        # ]

        return {
            "username": response["attributes"]["uid"][0],
            "password": response["attributes"]["userPassword"][0],
            "gid": response["attributes"]["gidNumber"],
            "mail": response["attributes"]["mail"][0],
            "mobile": response["attributes"]["telephoneNumber"][0] if response["attributes"]["telephoneNumber"] else [],
            # "locked": True if response['attributes']['pwdAccountLockedTime'] else False,
            # "description": '\n'.join(response['attributes']['description']),
        }

    def get_users(self, filter_syntax="(&(uid=*)(objectClass=posixAccount))"):
        attributes = ['uid', 'description', 'mail', 'gidNumber', 'telephoneNumber']

        conn = self.bind()
        conn.search(search_base=self.user_dn, search_filter=filter_syntax,
                    search_scope=ldap3.SUBTREE, attributes=attributes)

        groups = {str(group["id"]): group["description"] or group["name"] for group in self.get_groups()}
        users = [
            {
                "username": entry['attributes']['uid'][0],
                "mail": entry['attributes']['mail'][0],
                "mobile": entry["attributes"]["telephoneNumber"][0] if entry["attributes"]["telephoneNumber"] else '',
                "group": groups.get(str(entry['attributes']["gidNumber"]), str(entry['attributes']["gidNumber"])),
                "gid": entry['attributes']["gidNumber"],
                # "locked": True if entry['attributes']['pwdAccountLockedTime'] else False,
                "description": '\n'.join(entry['attributes']['description']),
            }

            for entry in conn.response
            if 'attributes' in entry and all(attr in entry['attributes'] for attr in attributes)
        ]

        conn.unbind()
        return users

    def get_user_group(self, uid):
        attributes = ['cn', 'gidNumber']

        conn = self.bind()
        conn.search(search_base=self.base_dn,
                    search_filter="(&(memberUid=%s)(objectClass=posixGroup))" % uid,
                    search_scope=ldap3.SUBTREE, attributes=attributes)

        try:
            response = conn.response[0]
        except IndexError:
            return None

        conn.unbind()
        return response['attributes']['gidNumber']

    def modify_user(self, uid, attributes):
        conn = self.bind()
        changes = {}
        for attr_name, attr_value in attributes.items():
            changes.update({
                attr_name: [(ldap3.MODIFY_REPLACE, attr_value if isinstance(attr_value, list) else [attr_value])]
            })

        conn.modify(dn=self.get_user_dn(uid=uid), changes=changes)
        conn.unbind()
        return conn.result

    def modify_group(self, gid, attributes):
        conn = self.bind()
        changes = {}
        for attr_name, attr_value in attributes.items():
            changes.update({
                attr_name: [(ldap3.MODIFY_REPLACE, attr_value if isinstance(attr_value, list) else [attr_value])]
            })

        conn.modify(dn=self.get_group_dn(gid=gid), changes=changes)
        conn.unbind()
        return conn.result

    def delete_user_from_group(self, uid, gid):
        conn = self.bind()

        group_dn = self.get_group_dn(gid)
        if not group_dn:
            return

        conn.modify(dn=group_dn, changes={"memberUid": [(ldap3.MODIFY_DELETE, uid)]})
        conn.unbind()

        return conn.result

    def add_user_to_group(self, uid, gid):
        conn = self.bind()

        group_dn = self.get_group_dn(gid)
        if not group_dn:
            return

        conn.modify(dn=group_dn, changes={"memberUid": [(ldap3.MODIFY_ADD, uid)]})
        conn.unbind()

        return conn.result

    def add_user_to_group_of_unique(self, uid, group_cn):
        conn = self.bind()

        user_dn = self.get_user_dn(uid)
        group_dn = self.get_group_of_unique_dn(group_cn)
        if not (group_dn or user_dn):
            return

        conn.modify(dn=group_dn, changes={"uniqueMember": [(ldap3.MODIFY_ADD, user_dn)]})
        conn.unbind()

        return conn.result

    def delete_user_from_group_of_unique(self, user_dn, group_cn):
        conn = self.bind()
        group_dn = self.get_group_of_unique_dn(group_cn)
        if not (group_dn or user_dn):
            return

        conn.modify(dn=group_dn, changes={"uniqueMember": [(ldap3.MODIFY_DELETE, user_dn)]})
        conn.unbind()

        return conn.result

    def delete_user(self, uid):
        conn = self.bind()
        user_dn = self.get_user_dn(uid=uid)

        if not user_dn:
            return

        conn.delete(user_dn)

        gid = self.get_user_group(uid=uid)
        if gid: self.delete_user_from_group(uid, gid)

        vpn_groups = [
            vpn_group["cn"]
            for vpn_group in self.get_group_of_unique_list(
                filter_syntax="(&(cn=*)(objectClass=groupOfUniqueNames)(uniqueMember=%s))" % user_dn
            )
        ]

        if vpn_groups:
            for vpn_group in vpn_groups: self.delete_user_from_group_of_unique(user_dn, vpn_group)

        conn.unbind()

        return conn.result

    def add_user(self, attributes, object_classes=None):
        if not object_classes:
            object_classes = ['shadowAccount', 'posixAccount', 'inetOrgPerson', 'pwdPolicy', 'top']

        conn = self.bind()
        conn.add(dn="uid={uid},{user_dn}".format(uid=attributes['uid'], user_dn=self.user_dn),
                 object_class=object_classes,
                 attributes=attributes)

        conn.unbind()
        self.add_user_to_group(attributes['uid'], attributes['gidNumber'])

        return conn.result

    def add_ou(self, attributes, object_classes=None):
        if not object_classes:
            object_classes = ['top', 'organizationalUnit']

        conn = self.bind()
        conn.add(dn="ou={ou_name},{base_dn}".format(
            ou_name=attributes['ou'],
            base_dn=self.base_dn
        ),
            object_class=object_classes,
            attributes=attributes)
        conn.unbind()
        return conn.result

    def add_dc(self, attributes, object_classes=None):
        if not object_classes:
            object_classes = ['top', 'organization', 'dcObject']

        conn = self.bind()
        conn.add(dn=self.base_dn,
                 object_class=object_classes,
                 attributes=attributes)
        conn.unbind()
        return conn.result

    def modify_password(self, uid, old_password, new_password):
        conn = self.bind_user(uid, old_password, get_connection=True)
        print(conn.bound)
        if not conn.bound:
            raise LDAPBindError("user %s can not be bound" % uid)

        user_dn = self.get_user_dn(uid=uid)
        print(user_dn)
        if not user_dn:
            return

        h_old_password = ldap_salted_sha1.hash(old_password.encode())
        print(h_old_password)
        h_new_password = ldap_salted_sha1.hash(new_password.encode())
        print(h_new_password)
        modification = ModifyPassword(conn, user_dn, h_old_password, h_new_password)
        modification.send()
        conn.unbind()

        return conn.result

    def reset_password(self, uid, new_password):
        user_dn = self.get_user_dn(uid=uid)

        conn = self.bind()
        h_new_password = ldap_salted_sha1.hash(new_password.encode())

        conn.modify(dn=user_dn, changes={"userPassword": [(ldap3.MODIFY_REPLACE, h_new_password)]})
        conn.unbind()

        return conn.result

    # def reset_password_by_hash(self, uid, password_hash):
    #     user_dn = self.get_user_dn(uid=uid)
    #
    #     conn = self.bind()
    #     h_new_password = password_hash
    #
    #     conn.modify(dn=user_dn, changes={"userPassword": [(ldap3.MODIFY_REPLACE, h_new_password)]})
    #     conn.unbind()
    #
    #     return conn.result


if __name__ == '__main__':
    _config = {
        'LDAP_HOST': '106.52.172.23',
        'LDAP_PORT': 389,
        'LDAP_PROTOCOL': 'ldap',
        'LDAP_USER': 'admin',
        'LDAP_AUTH': 'UhZGmWmQ8BQfRxZ',
        'LDAP_BASE_DN': 'dc=welab,dc=co',
        'LDAP_USER_DN': 'ou=Users,dc=welab,dc=co',
        'LDAP_GROUP_DN': 'ou=Groups,dc=welab,dc=co'
    }

    ol = AuthLdapHandler()
    # print(ol.bind_user(uid="boll.lai", password="123456789"))
    # print(ol.reset_password(uid="boll.lai", new_password="123456789"))
    # user = ol.get_user(uid="boll.lai")
    # print(to_str(user.get("password")))
    print(ol.get_users())
    # print(ol.modify_password(uid="boll.lai", old_password="1234567", new_password="1234567"))
    # print(ol.get_available_uid(gid=1001))
    # print(ol.get_user_uidNumber(uid="boll.lai"))
